import os
import re
import json
import time
import uuid
import pymysql
from fastapi import FastAPI, UploadFile, File, HTTPException, Query
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
import dashscope

# ========== FastAPI 应用初始化 ==========
app = FastAPI()

# 跨域支持
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 静态目录（用于展示发票图片）
UPLOAD_DIR = r"D:\VSCode\发票管理系统后端文件\uploads"
os.makedirs(UPLOAD_DIR, exist_ok=True)
app.mount("/static", StaticFiles(directory=os.path.dirname(UPLOAD_DIR)), name="static")

# ========== 配置 ==========
DB_CONFIG = {
    "host": "localhost",
    "user": "root",
    "password": "123456",
    "database": "invoice_manage",
    "charset": "utf8mb4",
    "autocommit": False
}
TABLE_NAME = "invoice"
OCR_API_KEY = os.getenv('OCR_API_KEY')
MAX_RETRIES = 3

# ========== 工具函数 ==========

def clean_text(text):
    text = text.replace('\u3000', ' ').replace('\u00a0', ' ')
    if text and text[0] == '\ufeff':
        text = text[1:]
    text = text.replace('￥', '')
    text = text.replace('\r\n', '\n').replace('\r', '\n')
    text = re.sub(r'[ \t]+', ' ', text)
    return text.strip()

def extract_json(text):
    """
    提取OCR返回结果中的第一个完整JSON对象，并校验必要字段
    """
    match = re.search(r'\{[\s\S]*\}', text)
    if not match:
        raise ValueError("未找到JSON结构")

    json_text = clean_text(match.group())
    try:
        parsed = json.loads(json_text)
    except json.JSONDecodeError as e:
        with open('error_json_debug.txt', 'w', encoding='utf-8') as f:
            f.write(json_text)
        raise ValueError("JSON解析失败")

    # 校验必要字段，防止写入空值导致数据库出错
    required_fields = ['invoice_code', 'issue_date', 'buyer_info', 'seller_info', 'goods_info', 'total_amount']
    for field in required_fields:
        if field not in parsed or not parsed[field]:
            raise ValueError(f"字段缺失或为空：{field}")

    return parsed

def format_date(date_str):
    if not date_str:
        return None
    return re.sub(r'[年月]', '-', date_str).replace('日', '').strip()

def call_ocr_with_retry(image_path, retries=MAX_RETRIES):
    messages = [
        {
            "role": "user",
            "content": [
                {"image": image_path},
                {"text": (
                    "请从这张发票中提取以下字段，并以标准 JSON 格式回答："
                    "invoice_code, issue_date, buyer_info, seller_info, goods_info, total_amount。"
                    "请不要添加任何解释说明。"
                )}
            ]
        }
    ]
    for attempt in range(retries):
        try:
            response = dashscope.MultiModalConversation.call(
                api_key=OCR_API_KEY,
                model='qwen-vl-max-latest',
                messages=messages
            )
            if response.status_code == 200:
                print(f"OCR成功（第{attempt + 1}次）")
                return response.output.choices[0].message.content[0]['text']
            else:
                print(f"OCR失败，错误码:{response.code} 信息:{response.message}")
        except Exception as e:
            print(f"第{attempt + 1}次调用异常：{e}")
        time.sleep(2)
    raise Exception("OCR请求失败，超过最大重试次数")

def insert_invoice_and_image(data, uuid_filename, original_filename, full_path):
    conn = pymysql.connect(**DB_CONFIG)
    try:
        with conn.cursor() as cursor:
            invoice_sql = f"""
                INSERT INTO {TABLE_NAME}
                (invoice_code, issue_date, buyer_info, seller_info, goods_info, total_amount)
                VALUES (%s, %s, %s, %s, %s, %s)
            """
            cursor.execute(invoice_sql, (
                data['invoice_code'],
                format_date(data['issue_date']),
                json.dumps(data['buyer_info'], ensure_ascii=False),
                json.dumps(data['seller_info'], ensure_ascii=False),
                json.dumps(data['goods_info'], ensure_ascii=False),
                data['total_amount']
            ))
            invoice_id = cursor.lastrowid
            print(f"发票插入成功，invoice_id={invoice_id}")

            image_sql = """
                INSERT INTO invoice_image (uuid_name, original_name, path, invoice_id)
                VALUES (%s, %s, %s, %s)
            """
            cursor.execute(image_sql, (
                uuid_filename,
                original_filename,
                full_path,
                invoice_id
            ))
            print(f"图片信息插入成功，文件名={uuid_filename}")

            conn.commit()
            return invoice_id
    except Exception as e:
        conn.rollback()
        print("数据库操作失败：", str(e))
        raise e
    finally:
        conn.close()

# ========== 上传接口 ==========
@app.post("/upload_invoice/")
async def upload_invoice(file: UploadFile = File(...)):
    if not file.filename.lower().endswith(('.png', '.jpg', '.jpeg')):
        raise HTTPException(status_code=400, detail="只支持图片格式")

    original_filename = file.filename
    file_uuid = str(uuid.uuid4())
    ext = os.path.splitext(file.filename)[-1].lower()
    uuid_filename = f"{file_uuid}{ext}"
    full_path = os.path.join(UPLOAD_DIR, uuid_filename)

    with open(full_path, "wb") as f:
        f.write(await file.read())

    print(f"文件保存成功，路径={full_path}")

    try:
        ocr_text = call_ocr_with_retry(full_path)
        parsed_data = extract_json(ocr_text)
        invoice_id = insert_invoice_and_image(parsed_data, uuid_filename, original_filename, full_path)
        return JSONResponse(content={
            "message": "识别并保存成功",
            "invoice_id": invoice_id,
            "image_file": uuid_filename,
            "data": parsed_data
        })
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# ========== 查询接口 ==========
@app.get("/search_invoices/")
def search_invoices(
    start_date: str = Query(None),
    end_date: str = Query(None),
    keyword: str = Query(None)
):
    conn = pymysql.connect(**DB_CONFIG)
    try:
        with conn.cursor(pymysql.cursors.DictCursor) as cursor:
            sql = f"""
                SELECT i.id, i.invoice_code, i.issue_date, i.buyer_info, i.seller_info, i.total_amount,
                       img.uuid_name AS image_file
                FROM {TABLE_NAME} i
                LEFT JOIN invoice_image img ON i.id = img.invoice_id
                WHERE 1=1
            """
            params = []
            if start_date:
                sql += " AND i.issue_date >= %s"
                params.append(start_date)
            if end_date:
                sql += " AND i.issue_date <= %s"
                params.append(end_date)
            if keyword:
                sql += " AND (i.invoice_code LIKE %s OR i.buyer_info LIKE %s OR i.seller_info LIKE %s)"
                keyword_pattern = f"%{keyword}%"
                params += [keyword_pattern] * 3

            cursor.execute(sql, params)
            rows = cursor.fetchall()
            for r in rows:
                r['buyer_info'] = json.loads(r['buyer_info'])
                r['seller_info'] = json.loads(r['seller_info'])
            return rows
    finally:
        conn.close()

# ========== 统计接口 ==========
@app.get("/statistics/")
def statistics(
    start_date: str = Query(None),
    end_date: str = Query(None),
    keyword: str = Query(None)
):
    conn = pymysql.connect(**DB_CONFIG)
    try:
        with conn.cursor() as cursor:
            sql = f"""
                SELECT COUNT(*) AS count, SUM(total_amount) AS total_amount
                FROM {TABLE_NAME}
                WHERE 1=1
            """
            params = []
            if start_date:
                sql += " AND issue_date >= %s"
                params.append(start_date)
            if end_date:
                sql += " AND issue_date <= %s"
                params.append(end_date)
            if keyword:
                sql += " AND (invoice_code LIKE %s OR buyer_info LIKE %s OR seller_info LIKE %s)"
                keyword_pattern = f"%{keyword}%"
                params += [keyword_pattern] * 3

            cursor.execute(sql, params)
            count, total = cursor.fetchone()
            return {"count": count or 0, "total_amount": float(total or 0)}
    finally:
        conn.close()

